package com.lee.limiter.maplimit;

import com.google.common.base.Preconditions;

/**
 * 默认为纯java实现的限流。
 * 有诸多弊端，如线程安全问题。
 * TokenBucketPool内存存储问题,面对海量的请求。pool的内容容量也增加到很大
 * 基于令牌桶的算法实现的limiter，以恒定的速率生成token。
 * User:liwei
 * Date:2017/9/1
 * Time:14:26
 */
public class RateLimiter {

    private RateLimitPolicy rateLimitPolicy;

    /**
     * 申请token权限
     *
     * @param identityKey
     * @return
     */
    public boolean access(String identityKey) {
        String tokenKey = rateLimitPolicy.getBucketKey(identityKey);
        TokenBucketPool pool = TokenBucketPool.getInstace();

        TokenBucket tokenBucket = pool.getBucketByKey(tokenKey);
        long currentTokensRemaining;//用于计算请求过来的时候当前的token
        long currentFillTime = System.currentTimeMillis();//当前时间
        if (tokenBucket == null) {
            synchronized (pool){
                if (pool.getBucketByKey(tokenKey) == null) {
                    currentTokensRemaining = rateLimitPolicy.getMaxBurstTokens();
                    tokenBucket = new TokenBucket(System.currentTimeMillis(), currentTokensRemaining - 1);
                    pool.putBucket(tokenKey, tokenBucket);
                }
            }
            return true;
        } else {
            synchronized (tokenBucket) {
                long lastRefillTime = tokenBucket.getLastRefillTime();//上次填充的时间
                long tokenRemaining = tokenBucket.getTokensRemaining();//桶中存活的token
                long intervalSinceLast = currentFillTime - lastRefillTime;

                if (intervalSinceLast > rateLimitPolicy.getIntervalInMills()) {
                    currentTokensRemaining = rateLimitPolicy.getMaxBurstTokens();
                    tokenBucket.setLastRefillTime(currentFillTime);
                } else {
                    //新生成的tokens总数
                    long grantedTokens = (intervalSinceLast / rateLimitPolicy.getIntervalPerPermit());
                    if (grantedTokens > 0) {
                        long padMillis = intervalSinceLast % rateLimitPolicy.getIntervalPerPermit();
                        tokenBucket.setLastRefillTime(currentFillTime - padMillis);
                    }
                    currentTokensRemaining = Math.min(grantedTokens + tokenRemaining, rateLimitPolicy.getTokenSize());
                }
                assert currentTokensRemaining >= 0;
                if (currentTokensRemaining == 0) {
                    tokenBucket.setTokensRemaining(currentTokensRemaining);
                    return false;
                } else {
                    tokenBucket.setTokensRemaining(currentTokensRemaining - 1);
                    return true;
                }
            }
        }
    }

    public static class Builder {

        public RateLimiter newRateLimiter(RateLimitPolicy policy) {
            Preconditions.checkNotNull(policy);
            RateLimiter limiter = new RateLimiter();
            limiter.rateLimitPolicy = policy;
            return limiter;
        }
    }
}
